use std::collections::HashMap;

/// 单变量线性拟合结果
#[derive(Debug, Clone)]
pub struct LinearFitResult {
    pub slope: f64,
    pub intercept: f64,
    pub r2: f64,
}

/// 单变量线性拟合
pub fn single_linear(y: &[f64], x: Option<&[f64]>) -> LinearFitResult {
    let x_data = match x {
        Some(x) => x.to_vec(),
        None => (0..y.len()).map(|i| i as f64).collect(),
    };
    
    let n = x_data.len() as f64;
    let x_sum: f64 = x_data.iter().sum();
    let y_sum: f64 = y.iter().sum();
    let x_squared_sum: f64 = x_data.iter().map(|x| x * x).sum();
    let xy_sum: f64 = x_data.iter().zip(y.iter()).map(|(x, y)| x * y).sum();
    
    let delta = n * x_squared_sum - x_sum * x_sum;
    
    if delta.abs() < 1e-10 {
        return LinearFitResult {
            slope: 0.0,
            intercept: 0.0,
            r2: 0.0,
        };
    }
    
    let intercept = (x_squared_sum * y_sum - x_sum * xy_sum) / delta;
    let slope = (n * xy_sum - x_sum * y_sum) / delta;
    
    // 计算R²
    let y_mean = y_sum / n;
    let ss_tot: f64 = y.iter().map(|yi| (yi - y_mean).powi(2)).sum();
    let ss_err: f64 = x_data
        .iter()
        .zip(y.iter())
        .map(|(xi, yi)| (yi - slope * xi - intercept).powi(2))
        .sum();
    
    let r2 = if ss_tot.abs() < 1e-10 {
        0.0
    } else {
        1.0 - ss_err / ss_tot
    };
    
    LinearFitResult {
        slope: (slope * 10000.0).round() / 10000.0,
        intercept: (intercept * 10000.0).round() / 10000.0,
        r2: (r2 * 10000.0).round() / 10000.0,
    }
}

/// 计算皮尔逊相关系数
pub fn pearson_correlation(x: &[f64], y: &[f64]) -> f64 {
    if x.len() != y.len() || x.is_empty() {
        return 0.0;
    }
    
    let n = x.len() as f64;
    let x_sum: f64 = x.iter().sum();
    let y_sum: f64 = y.iter().sum();
    let xy_sum: f64 = x.iter().zip(y.iter()).map(|(x, y)| x * y).sum();
    let x_squared_sum: f64 = x.iter().map(|x| x * x).sum();
    let y_squared_sum: f64 = y.iter().map(|y| y * y).sum();
    
    let numerator = n * xy_sum - x_sum * y_sum;
    let denominator = ((n * x_squared_sum - x_sum * x_sum) * (n * y_squared_sum - y_sum * y_sum)).sqrt();
    
    if denominator.abs() < 1e-10 {
        0.0
    } else {
        numerator / denominator
    }
}

/// 计算斯皮尔曼相关系数
pub fn spearman_correlation(x: &[f64], y: &[f64]) -> f64 {
    if x.len() != y.len() || x.is_empty() {
        return 0.0;
    }
    
    // 计算排名
    let x_ranks = rank_data(x);
    let y_ranks = rank_data(y);
    
    // 使用排名计算皮尔逊相关系数
    pearson_correlation(&x_ranks, &y_ranks)
}

/// 为数据计算排名
fn rank_data(data: &[f64]) -> Vec<f64> {
    let mut indexed_data: Vec<(usize, f64)> = data.iter().enumerate().map(|(i, &v)| (i, v)).collect();
    indexed_data.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
    
    let mut ranks = vec![0.0; data.len()];
    let mut i = 0;
    while i < indexed_data.len() {
        let mut j = i;
        let current_value = indexed_data[i].1;
        
        // 找到所有相同值的索引
        while j < indexed_data.len() && indexed_data[j].1 == current_value {
            j += 1;
        }
        
        // 计算平均排名
        let avg_rank = (i + j - 1) as f64 / 2.0 + 1.0;
        for k in i..j {
            ranks[indexed_data[k].0] = avg_rank;
        }
        
        i = j;
    }
    
    ranks
}

/// 计算肯德尔相关系数
pub fn kendall_correlation(x: &[f64], y: &[f64]) -> f64 {
    if x.len() != y.len() || x.len() < 2 {
        return 0.0;
    }
    
    let mut concordant = 0;
    let mut discordant = 0;
    
    for i in 0..x.len() {
        for j in (i + 1)..x.len() {
            let x_diff = x[j] - x[i];
            let y_diff = y[j] - y[i];
            
            if x_diff * y_diff > 0.0 {
                concordant += 1;
            } else if x_diff * y_diff < 0.0 {
                discordant += 1;
            }
            // 如果x_diff或y_diff为0，不计入统计
        }
    }
    
    let total = concordant + discordant;
    if total == 0 {
        0.0
    } else {
        (concordant as f64 - discordant as f64) / total as f64
    }
}

/// 计算IC值（信息系数）
pub fn calculate_ic(x: &[f64], y: &[f64], method: &str) -> f64 {
    match method.to_lowercase().as_str() {
        "pearson" => pearson_correlation(x, y),
        "spearman" => spearman_correlation(x, y),
        "kendall" => kendall_correlation(x, y),
        _ => pearson_correlation(x, y), // 默认使用皮尔逊相关系数
    }
}

/// 计算滚动相关系数
pub fn rolling_correlation(
    x: &[f64],
    y: &[f64],
    window: usize,
    method: &str,
) -> Vec<f64> {
    if x.len() != y.len() || window > x.len() {
        return vec![];
    }
    
    let mut correlations = Vec::new();
    
    for i in (window - 1)..x.len() {
        let x_window = &x[(i - window + 1)..=i];
        let y_window = &y[(i - window + 1)..=i];
        
        let correlation = calculate_ic(x_window, y_window, method);
        correlations.push(correlation);
    }
    
    correlations
}

/// 计算相关性矩阵
pub fn correlation_matrix(data: &[Vec<f64>], method: &str) -> Vec<Vec<f64>> {
    let n = data.len();
    let mut matrix = vec![vec![0.0; n]; n];
    
    for i in 0..n {
        for j in 0..n {
            if i == j {
                matrix[i][j] = 1.0;
            } else {
                matrix[i][j] = calculate_ic(&data[i], &data[j], method);
            }
        }
    }
    
    matrix
}

/// 计算IC统计信息
pub fn ic_statistics(ic_values: &[f64]) -> HashMap<String, f64> {
    let mut stats = HashMap::new();
    
    if ic_values.is_empty() {
        return stats;
    }
    
    let mean = ic_values.iter().sum::<f64>() / ic_values.len() as f64;
    let variance = ic_values
        .iter()
        .map(|x| (x - mean).powi(2))
        .sum::<f64>()
        / ic_values.len() as f64;
    let std_dev = variance.sqrt();
    
    let positive_count = ic_values.iter().filter(|&&x| x > 0.0).count();
    let hit_rate = positive_count as f64 / ic_values.len() as f64;
    
    stats.insert("mean".to_string(), mean);
    stats.insert("std_dev".to_string(), std_dev);
    stats.insert("hit_rate".to_string(), hit_rate);
    stats.insert("min".to_string(), *ic_values.iter().min_by(|a, b| a.partial_cmp(b).unwrap()).unwrap());
    stats.insert("max".to_string(), *ic_values.iter().max_by(|a, b| a.partial_cmp(b).unwrap()).unwrap());
    
    stats
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_single_linear() {
        let y = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let result = single_linear(&y, None);
        
        assert!((result.slope - 1.0).abs() < 0.01);
        assert!((result.r2 - 1.0).abs() < 0.01);
    }
    
    #[test]
    fn test_pearson_correlation() {
        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let y = vec![2.0, 4.0, 6.0, 8.0, 10.0];
        
        let correlation = pearson_correlation(&x, &y);
        assert!((correlation - 1.0).abs() < 0.01);
    }
    
    #[test]
    fn test_spearman_correlation() {
        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let y = vec![2.0, 4.0, 6.0, 8.0, 10.0];
        
        let correlation = spearman_correlation(&x, &y);
        assert!((correlation - 1.0).abs() < 0.01);
    }
} 